import cv2
import numpy as np
from matplotlib import pyplot as plt
import time

from PIL import ImageGrab

time.sleep(5)
# ImageGrab.grab().convert('RGB').show()
img = np.array(ImageGrab.grab().convert('RGB'))
# img = np.array(ImageGrab.grab().convert('L'))
template = cv2.imread('images/mb.png', 1)
img2 = img.copy()
w, h ,k= template.shape[::-1]

# All the 6 methods for comparison in a list
methods = ['cv2.TM_CCOEFF', 'cv2.TM_CCOEFF_NORMED', 'cv2.TM_CCORR',
           'cv2.TM_CCORR_NORMED', 'cv2.TM_SQDIFF', 'cv2.TM_SQDIFF_NORMED']

methods = ['cv2.TM_CCOEFF', ]
for meth in methods:
    img = img2.copy()
    method = eval(meth)

    # 匹配
    res = cv2.matchTemplate(img, template, method)# 小图块 ，大图块   ，查找模式
    min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)

    # If the method is TM_SQDIFF or TM_SQDIFF_NORMED, take minimum
    if method in [cv2.TM_SQDIFF, cv2.TM_SQDIFF_NORMED]:
        top_left = min_loc
    else:
        top_left = max_loc

    bottom_right = (top_left[0] + w, top_left[1] + h) #计算偏移

    cv2.rectangle(img, top_left, bottom_right, 255, 100)#矩形

    plt.subplot(121), plt.imshow(res, cmap='gray')
    plt.title('Matching Result'), plt.xticks([]), plt.yticks([]) #匹配结果
    plt.subplot(122), plt.imshow(img, cmap='gray')
    plt.title('Detected Point'), plt.xticks([]), plt.yticks([])
    plt.suptitle(meth)

    plt.show()